import glob
import gzip
import json
import mimetypes
import os
import re
import shutil
import tarfile
import tempfile
import zipfile
from urllib.parse import urlparse

import requests

from common.wrapper.orochi_wrapper import cost_time
from utils import config

_LOCAL_PREFIX = "file://"
_URI_RE = "https?://(.+)/(.+)"
_HTTP_PREFIX = "http(s)://"
_HEADERS_SUFFIX = "-headers"


@cost_time
def fetch_cache_model(filename=""):
    fetch_cache_uri(config.storage_uri(), filename)

@cost_time
def fetch_cache_uri(uri, filename=""):
    if filename == "":
        return download(uri, os.path.dirname(os.path.dirname(__file__)) + "/tmp/" + config.model_name(), False)
    return download(uri, os.path.dirname(os.path.dirname(__file__)) + "/tmp/" + config.model_name(),
                            False) + "/" + filename

def download(uri: str, out_dir: str = None, replace_if_present=True) -> str:
    print("Copying contents of " + uri + " to local")
    is_local = False
    if uri.startswith(_LOCAL_PREFIX) or os.path.exists(uri):
        is_local = True

    if out_dir is None:
        if is_local:
            # noop if out_dir is not set and the path is local
            return download_local(uri)
        out_dir = tempfile.mkdtemp()
    elif not os.path.exists(out_dir):
        os.makedirs(out_dir)

    if is_local:
        return download_local(uri, out_dir, replace_if_present)
    elif re.search(_URI_RE, uri):
        return download_from_uri(uri, out_dir, replace_if_present)
    else:
        raise Exception("Cannot recognize storage type for " + uri +
                        "\n'%s', '%s' are the current available storage type." %
                        (_LOCAL_PREFIX, _HTTP_PREFIX))


def download_local(uri, out_dir=None, replace_if_present=True):
    local_path = uri.replace(_LOCAL_PREFIX, "", 1)
    if not os.path.exists(local_path):
        raise RuntimeError("Local path %s does not exist." % uri)

    if out_dir is None:
        return local_path
    elif not os.path.isdir(out_dir):
        os.makedirs(out_dir)

    if os.path.isdir(local_path):
        local_path = os.path.join(local_path, "*")

    count = 0
    for src in glob.glob(local_path):
        _, tail = os.path.split(src)
        dest_path = os.path.join(out_dir, tail)
        if not replace_if_present and os.path.exists(dest_path):
            return out_dir
        print("Linking: " + src + " to " + dest_path)
        os.symlink(src, dest_path)
        count = count + 1

    # Unpack compressed file, supports .tgz, tar.gz and zip file formats.
    if count == 1:
        mimetype, _ = mimetypes.guess_type(dest_path)
        if mimetype in ["application/x-tar", "application/zip", "application/x-zip-compressed"]:
            unpack_archive_file(dest_path, mimetype, out_dir)

    return out_dir


def download_from_uri(uri, out_dir=None, replace_if_present=True):
    url = urlparse(uri)
    filename = os.path.basename(url.path)
    mimetype, encoding = mimetypes.guess_type(url.path)
    local_path = os.path.join(out_dir, filename)

    if not replace_if_present and os.path.exists(local_path):
        return out_dir

    if filename == '':
        raise ValueError('No filename contained in URI: %s' % uri)

    host_uri = url.hostname

    headers_json = os.getenv(host_uri + _HEADERS_SUFFIX, "{}")
    headers = json.loads(headers_json)

    with requests.get(uri, stream=True, headers=headers) as response:
        if response.status_code != 200:
            raise RuntimeError("URI: %s returned a %s response code." % (uri, response.status_code))
        zip_content_types = (
            'application/x-zip-compressed', 'application/zip', 'application/zip-compressed', 'application/octet-stream')
        if mimetype == 'application/zip' and not response.headers.get('Content-Type', '').startswith(zip_content_types):
            raise RuntimeError("URI: %s did not respond with any of following \'Content-Type\': " % uri +
                               ", ".join(zip_content_types))
        tar_content_types = ('application/x-tar', 'application/x-gtar', 'application/x-gzip', 'application/gzip')
        if mimetype == 'application/x-tar' and not response.headers.get('Content-Type', '') \
                .startswith(tar_content_types):
            raise RuntimeError("URI: %s did not respond with any of following \'Content-Type\': " % uri +
                               ", ".join(tar_content_types))
        if (
                mimetype != 'application/zip' and mimetype != 'application/x-tar'
                and mimetype != 'application/x-zip-compressed') and not response.headers.get(
            'Content-Type', '').startswith('application/octet-stream'):
            raise RuntimeError("URI: %s did not respond with \'Content-Type\': \'application/octet-stream\'"
                               % uri)

        if encoding == 'gzip':
            stream = gzip.GzipFile(fileobj=response.raw)
            local_path = os.path.join(out_dir, filename + '.tar')
        else:
            stream = response.raw
        with open(local_path, 'wb') as out:
            shutil.copyfileobj(stream, out)

    if mimetype in ["application/x-tar", "application/zip", "application/x-zip-compressed", 'application/octet-stream']:
        unpack_archive_file(local_path, mimetype, out_dir)

    return out_dir


def unpack_archive_file(file_path, mimetype, target_dir=None):
    if not target_dir:
        target_dir = os.path.dirname(file_path)

    try:
        print("Unpacking: %s" % file_path)
        if mimetype == "application/x-tar":
            archive = tarfile.open(file_path, 'r', encoding='utf-8')
        else:
            archive = zipfile.ZipFile(file_path, 'r')
        archive.extractall(target_dir)
        archive.close()
    except (tarfile.TarError, zipfile.BadZipfile):
        raise RuntimeError("Failed to unpack archieve file. The file format is not valid.")
